﻿#define PARALLEL

using System;
using System.Collections.Generic;
using System.Text;
using System.Net;
using System.IO;
using System.Xml;
using System.Web.Script.Serialization;

namespace KrisThom
{
    public class VirusTotalScan
    {

        public string VTApiKey { get; set; }

        public VirusTotalScan(String apikey)
        {
            VTApiKey = apikey;
        }

        /// <summary>
        /// Upload file to the service and output report link
        /// </summary>
        /// <param name="apikey"></param>
        /// <param name="uploadfile"></param>
        /// <param name="pageSource"></param>
        /// <param name="title"></param>
        /// <returns></returns>
        public static VirusTotalResults Upload(String apikey, FileInfo uploadfile)
        {
            try
            {
                List<Byte> byteList = new List<byte>();

                string boundary = "----------" + DateTime.Now.Ticks.ToString();

                String header =
                    "--" + boundary +
                    "\r\n" +
                    "Content-Disposition: form-data; name=\"key\"" +
                    "\r\n\r\n" +
                    apikey +
                    "\r\n" +
                    "--" + boundary +
                    "\r\n" +
                    "Content-Disposition: form-data; name=\"file\"; filename=\""+ uploadfile.Name + "\"" +
                    "\r\n" +
                    "Content-Type: application/octet-stream" +
                    "\r\n\r\n";

                byte[] headerbyte = Encoding.ASCII.GetBytes(header);

                byte[] filebyte = File.ReadAllBytes(uploadfile.FullName);
                
                byte[] trailer = Encoding.ASCII.GetBytes("\r\n--" + boundary + "--\r\n");

                byteList.AddRange(headerbyte);
                byteList.AddRange(filebyte);
                byteList.AddRange(trailer);

                byte[] data = byteList.ToArray();


                HttpWebRequest request = (HttpWebRequest)WebRequest.Create("http://www.virustotal.com/api/scan_file.json");
                request.Method = "POST";
                //request.Proxy = new WebProxy("Http://www.virustotal.com");
                request.ProtocolVersion = HttpVersion.Version10;
                request.ContentType = "multipart/form-data; boundary=" + boundary;
                request.ContentLength = data.Length;
                request.KeepAlive = false;
                ServicePointManager.Expect100Continue = false;
     
                // execute the request
                Stream dataStream = request.GetRequestStream();
                dataStream.Write(data, 0, data.Length);
                dataStream.Close();

                WebResponse response = request.GetResponse();
                using (HttpWebResponse HttpResponse = (HttpWebResponse)request.GetResponse())
                {
                    // we will read data via the response stream
                    using (Stream resStream = response.GetResponseStream())
                    {
                        String rawVT = new StreamReader(resStream).ReadToEnd();

                        return new VirusTotalResults(rawVT);
                    }
                }
            }
            catch
            {
                throw;
            }
        }

        /// <summary>
        /// Scan multiple files, this is a blocking call, maximum 20 requests every 5 minutes
        /// </summary>
        /// <param name="apikey"></param>
        /// <param name="uploadfiles"></param>
        /// <returns></returns>
        public VirusTotalResults[] Run(FileInfo[] uploadfiles)
        {
            int pauseminutes = 5;
            int pausems = (int)TimeSpan.FromMinutes(pauseminutes).TotalMilliseconds;
            List<VirusTotalResults> vtlist = new List<VirusTotalResults>();
            try
            {
                List<FileInfo> files = new List<FileInfo>();
                files.AddRange(uploadfiles);

                List<FileInfo>[] batchlist = Helper.Partition(files, 3);

                foreach (List<FileInfo> batch in batchlist)
                {
                    RunInternal(batch.ToArray());
                    System.Threading.Thread.Sleep(pausems);
                }
            }
            catch
            {
                throw;
            }
            
            return vtlist.ToArray();
        }

        private VirusTotalResults[] RunInternal(FileInfo[] uploadfiles)
        {
            try
            {
                List<VirusTotalResults> vtlist = new List<VirusTotalResults>();

#if (PARALLEL)
                lock (vtlist)
                {
                    System.Threading.Tasks.Parallel.ForEach<FileInfo>(uploadfiles, delegate(FileInfo file)
                    {
                        //VirusTotalResults vtresult = Upload(VTApiKey, file);
                        //vtlist.Add(vtresult);   
                    }
                    );
                }

#else
                foreach (FileInfo file in uploadfiles)
                {
                    VirusTotalResults vtresult = Upload(VTApiKey, file);
                    vtlist.Add(vtresult);                        
                }
#endif

                return vtlist.ToArray();
            }
            catch
            {
                throw;
            }
        }

        public static VirusTotalResults MD5(String apikey, string hash)
        {
            try
            {

                String postdata = String.Format("resource={0}&key={1}", hash, apikey);
                byte[] data = UTF8Encoding.UTF8.GetBytes(postdata);

                WebRequest request = WebRequest.Create("http://www.virustotal.com/api/get_file_report.json");
                request.Method = "POST";
                request.ContentType = "application/x-www-form-urlencoded";
                request.ContentLength = data.Length;
                request.Headers.Add("Accept-Encoding", "identity");
                ServicePointManager.Expect100Continue = false;

                // execute the request
                Stream dataStream = request.GetRequestStream();
                dataStream.Write(data, 0, data.Length);
                dataStream.Close();

                WebResponse response = request.GetResponse();
                using (HttpWebResponse HttpResponse = (HttpWebResponse)request.GetResponse())
                {
                    // we will read data via the response stream
                    using (Stream resStream = response.GetResponseStream())
                    {
                        String rawVT = new StreamReader(resStream).ReadToEnd();

                        return new VirusTotalResults(rawVT);
                    }
                }
            }
            catch
            {
                throw;
            }
        }

        #region Previous code before API became available
        
        //private static Boolean IsMalware(String input)
        //{
        //    Boolean bad = false;
        //    try
        //    {
        //        if (input.Contains("No results found for that hash."))
        //        {
        //            bad = false;
        //        }
        //        else if (input.Contains("positivo"))
        //        {
        //            bad = true;
        //        }
        //    }
        //    catch (Exception ex)
        //    {
        //        throw ex;
        //    }
        //    return bad;
        //}



        //public static Boolean MD5(String hash, out String pageSource, out String title)
        //{
        //    // used to build entire input
        //    StringBuilder sb = new StringBuilder();
        //    Boolean Malware = false;
        //    pageSource = String.Empty;
        //    title = String.Empty;

        //    try
        //    {
        //        //post data

        //        String postdata = String.Format("hash={0}&x=85&y=22",hash);

        //        ASCIIEncoding encoding = new ASCIIEncoding();
        //        byte[] data = encoding.GetBytes(postdata);


        //        // used on each read operation
        //        byte[] buf = new byte[8192];

        //        // prepare the web page we will be asking for
        //        HttpWebRequest request = (HttpWebRequest)WebRequest.Create("http://www.virustotal.com/vt/en/consultamd5");
        //        request.Method = "POST";
        //        request.ContentType = "application/x-www-form-urlencoded";
        //        request.ContentLength = data.Length;

        //        // execute the request
        //        Stream newStream = request.GetRequestStream();
        //        newStream.Write(data, 0, data.Length);

        //        HttpWebResponse response = (HttpWebResponse)request.GetResponse();

        //        // we will read data via the response stream
        //        Stream resStream = response.GetResponseStream();

        //        string tempString = null;
        //        int count = 0;

        //        do
        //        {
        //            // fill the buffer with data
        //            count = resStream.Read(buf, 0, buf.Length);

        //            // make sure we read some data
        //            if (count != 0)
        //            {
        //                // translate from bytes to ASCII text
        //                tempString = Encoding.ASCII.GetString(buf, 0, count);

        //                // continue building the string
        //                sb.Append(tempString);
        //            }
        //        }
        //        while (count > 0); // any more data to read?


        //        //check if malware was found
        //        Malware = IsMalware(sb.ToString());

        //        if (Malware)
        //        {
        //            title = ParseTitle(sb.ToString());
        //        }

        //        pageSource = sb.ToString();

        //    }
        //    catch (Exception ex)
        //    {
        //        throw ex;
        //    }
                
        //    // print out page source
        //    return Malware;
        //}

        //private static String ParseTitle(String input)
        //{
        //    //<title>Virustotal. MD5: 27b002ee170c751d14e030dacbb52b9f Heuristic.ADH Artemis!27B002EE170C Rogue:W32/AntiVirusElite.C </title>

        //    String output = String.Empty;

        //    try
        //    {
        //        foreach (string s in input.Split('\n'))
        //        {
        //            if (s.Contains("<title>"))
        //            {
        //                output = s;
        //                output = output.Replace("<title>", "");
        //                output = output.Replace("</title>", "");
        //                output = output.Substring(50).Trim();
        //            }
        //        }
        //    }
        //    catch (Exception ex)
        //    {
        //        throw ex;
        //    }

        //    return output;
        //}

        #endregion

    }
}
